{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "491768c4-ef73-479b-a4dc-8b792ee56eec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"RERUN_NOTEBOOK_ASSET\"]=\"serve-local\"\n",
    "\n",
    "import time\n",
    "import os\n",
    "import requests\n",
    "import cv2\n",
    "\n",
    "import pyarrow as pa\n",
    "\n",
    "import rerun as rr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32d011a3-cf96-4a19-a26e-de7c1b6fb629",
   "metadata": {},
   "outputs": [],
   "source": [
    "def download(url: str):\n",
    "    if not os.path.exists(\"datasets\"):\n",
    "        os.makedirs(\"datasets\")\n",
    "    filename = url.split(\"/\")[-1]\n",
    "    dest = f\"datasets/{filename}\"\n",
    "    if not os.path.exists(dest):\n",
    "        with open(dest, \"wb\") as f:\n",
    "            f.write(requests.get(url).content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "212bf808-b385-4a40-abe4-c107820283c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# OPF\n",
    "download(\"http://app.rerun.io/version/0.18.2/examples/open_photogrammetry_format.rrd\")\n",
    "opf = rr.dataframe.load_recording(\"datasets/open_photogrammetry_format.rrd\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37199b15-fdec-417e-bdfd-aa29bc720196",
   "metadata": {},
   "outputs": [],
   "source": [
    "# What timelines do we have?\n",
    "opf.schema().index_columns()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0713ffc6-cd4d-41e3-bf13-907192c52403",
   "metadata": {},
   "outputs": [],
   "source": [
    "# What components do we have?\n",
    "opf.schema().component_columns()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64d2f495-b6dd-473d-8857-2b5f64bc68e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract a single point-cloud\n",
    "\n",
    "cloud = rr.dataframe.ComponentColumnSelector(\"world/points\", rr.components.Position3D)\n",
    "\n",
    "batches = opf.view(index=\"image\", contents=\"world/points\").select([cloud])\n",
    "\n",
    "# Should the python APIs do this for us?\n",
    "table = pa.Table.from_batches(batches)\n",
    "\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8773b66d-1db8-4b92-8e9e-9bc121c2fa61",
   "metadata": {},
   "outputs": [],
   "source": [
    "cloud = table[0][0]\n",
    "type(cloud)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a420e92-a7c3-4185-9772-9e3368cb1d26",
   "metadata": {},
   "outputs": [],
   "source": [
    "points_np = cloud.values.flatten().to_numpy().reshape(-1, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6316465b-291b-4130-95f8-ded71b00d6be",
   "metadata": {},
   "outputs": [],
   "source": [
    "points_np[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2265ab54-c9a6-4dff-ac60-b4d4246ff370",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find the positions of all the camera estimates\n",
    "\n",
    "img_num = rr.dataframe.TimeColumnSelector(\"image\")\n",
    "logged = rr.dataframe.TimeColumnSelector(\"log_time\")\n",
    "\n",
    "pos = rr.dataframe.ComponentColumnSelector(\"world/cameras\", rr.components.Translation3D)\n",
    "quat = rr.dataframe.ComponentColumnSelector(\"world/cameras\", rr.components.TransformMat3x3)\n",
    "\n",
    "view = opf.view(index=\"image\", contents=\"world/cameras\")\n",
    "batches = view.select([img_num, logged, pos, quat])\n",
    "\n",
    "table = pa.Table.from_batches(batches)\n",
    "\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22beda5e-96fa-48ca-b226-8480117e8a07",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Annoying -- row traversal even though they were logged together :-(\n",
    "table = table.combine_chunks()\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa7c4049-44f5-44b6-9be7-16cb2006898e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Look at some samples (Polars would make this nicer)\n",
    "table.take([0, 10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fce44b1b-58fd-486d-8323-b1b448b2eacd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Let's feed this into a geometry library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "862c9440-733b-4f64-9488-57eb053973ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Maybe we only wanted to query poses for a few images\n",
    "# Find the positions of all the camera estimates\n",
    "\n",
    "batches = view.filter_range_sequence(start=10, end=20).select([img_num, logged, pos, quat])\n",
    "\n",
    "table = pa.Table.from_batches(batches)\n",
    "\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d9edee4-d6c2-4c6a-8074-acf7733cbbee",
   "metadata": {},
   "outputs": [],
   "source": [
    "table.take([0, 10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd185db4-f8f9-4332-8359-f509652e0301",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: latest-at"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8e52744-c138-437d-bc83-17eb17c3775d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's grab an image\n",
    "\n",
    "view = opf.view(index=\"image\", contents=\"/**\").filter_range_sequence(42,42)\n",
    "\n",
    "img_num = rr.dataframe.TimeColumnSelector(\"image\")\n",
    "\n",
    "# Getting these right is annoying -- would be nice to support '/**' somehow\n",
    "blob = rr.dataframe.ComponentColumnSelector(\"world/cameras/image\", rr.components.Blob)\n",
    "media = rr.dataframe.ComponentColumnSelector(\"world/cameras/image\", rr.components.MediaType)\n",
    "\n",
    "batches = view.select([blob, media])\n",
    "\n",
    "table = pa.Table.from_batches(batches)\n",
    "\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55a7d30e-f5b2-4b9a-98bd-0f272a37b18d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ^^ when we typo something, we just find no data\n",
    "table.schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1495ce3e-022c-4776-ba2a-2156f3527e92",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's grab an image for real this time\n",
    "\n",
    "# Getting these right is annoying -- would be nice to support '/**' somehow\n",
    "blob = rr.dataframe.ComponentColumnSelector(\"world/cameras/image/rgb\", rr.components.Blob)\n",
    "media = rr.dataframe.ComponentColumnSelector(\"world/cameras/image/rgb\", rr.components.MediaType)\n",
    "\n",
    "batches = view.select([blob, media])\n",
    "\n",
    "table = pa.Table.from_batches(batches)\n",
    "\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")\n",
    "\n",
    "## NOTE: Kind of surprising to find 2 rows here. The first one gets included because the VIEW contains\n",
    "## Contains static data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c997303-c515-44be-8522-1bd039491798",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Image is encoded as: {table[1][1][0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d9574c0-a177-48a0-a184-9b61cd6da76d",
   "metadata": {},
   "outputs": [],
   "source": [
    "image = image = cv2.imdecode(table[0][1][0].values.to_numpy(), cv2.IMREAD_COLOR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4127cfc4-6ff2-4136-85ef-77f084fddcee",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"preview\")\n",
    "\n",
    "rr.log(\"img\", rr.Image(image, color_model=\"bgr\"))\n",
    "\n",
    "rr.notebook_show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2e6bc2e-f202-446d-923c-72fb219a460c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# BGR shader not working on web in firefox :sob:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a48cacc0-cf92-43a5-a4c2-85382d69c591",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"preview\")\n",
    "\n",
    "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "rr.log(\"img\", rr.Image(image))\n",
    "\n",
    "rr.notebook_show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aa4ca6d-f3dd-4024-96b2-dcfe8abe4d1f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
